import time
from loguru import logger
import time
from typing import List, Union
import requests
import datetime
from loguru import logger
from io import BytesIO
from pydantic import BaseModel
from requests_toolbelt.multipart.encoder import MultipartEncoder

class MaaS:
    def __init__(self,
                 endpoint="http://maas.capitalonline.net",
                 token=""
                 ):
        self.endpoint = endpoint[0:-1] if endpoint[-1] == "/" else endpoint
        self.headers = {"Authorization": f"Token {token}", "Content-Type": "application/json"}


    def upload_file(self, filename):
        url = self.endpoint + "/api/files/UploadOneFileView/"
        with open(filename, "rb") as f:
            headers = {**self.headers}
            multipart_data = MultipartEncoder(
                fields={
                    'file': (filename, f, 'text/plain')
                }
            )
            headers['Content-Type'] = multipart_data.content_type
            headers['accept'] = 'application/json'
            try:
                resp = requests.post(url, data=multipart_data, headers=headers)
                resp.encoding = "utf8"
                if resp.status_code < 300: 
                    return resp.json().get("download_url")
                else:
                    logger.warning(f"上传文件失败: {resp.text}")
            except Exception as e:
                logger.exception(str(e))
                return ""
            return ""

    # 推理接口
    def infer(self, init_image, mask_image, image_num):
        url = self.endpoint + "/api/v1/image/creation/"
        data = {
            "parameters": {
                "function": "controlnet_inpaint",
                "timestamp": "2024-01-22 15:23:29.451",
                "parameters": {
                    "model": {
                        "id": 1651,
                        "name": "majicMIX-realistic-v6（麦橘写实）"
                    },
                    "controlnet": "inpaint",
                    "style": [],
                    "resolution": "680x512",
                    "strength": 0.5,
                    "condition_scale": 0.5,
                    "guidance_scale": 0.5,
                    "image_num": image_num,
                    "init_image": init_image,
                    "mask_image": mask_image,
                    "prompt": {
                        "positive": "1person",
                        "negative": ""
                    },
                    "seed": int(time.time()),
                    "denoising_steps": 20,
                    "scheduler": "ddpm",
                    "guidance_scale": 7.5
                }
            }
        }
        logger.info(f"绘图参数: {data}")
        return requests.post(url, json=data, headers=self.headers)

    def get_result(self, _id):
        url = self.endpoint + f"/api/v1/generation/task/status/{_id}/"
        return requests.get(url, headers=self.headers)


    def infer_with_result(self, init_image, mask_image,image_num, timeout=120):
        resp = self.infer(init_image=init_image, mask_image=mask_image, image_num=image_num)
        resp.encoding = "utf8"
        images_urls = []
        if resp.status_code < 300:
            request_id = resp.json().get("request_id")
            start_time = time.time()
            while True:
                result = self.get_result(_id=request_id)
                images_urls = []
                logger.info(f"绘图中... {result.json()}")
                if result.status_code < 300:
                    images = result.json().get("images", [])
                    if result.json().get("task_info", {}).get("status") == "completed" and images:
                        for i in images:
                            images_urls.append(i.get("image_uri"))
                    if len(images_urls) == image_num:
                        break
                else:
                    logger.warning(f"创作异常 {resp} {resp.text}")
                if time.time() - start_time > timeout:
                    break
                time.sleep(2)
        else:
            logger.warning(f"创作异常 {resp} {resp.text}")
        return images_urls
